import socket
import struct
import zlib
import threading
import tkinter as tk
from io import BytesIO
from PIL import Image, ImageTk
import time
import tkinter.messagebox as msgb
import tkinter.simpledialog as strq
import pyaudio  # 需安装：pip install pyaudio


class RemoteClient:
    def __init__(self, host='127.0.0.1', port=65432):
        self.host = host
        self.port = port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.running = False
        self.latest_frame = None  # 最新视频帧
        self.frame_lock = threading.Lock()  # 帧数据锁
        self.input_thread = None
        self.fps = 0
        self.last_frame_time = 0

        # 音频播放配置
        self.audio = pyaudio.PyAudio()
        self.audio_stream = self.audio.open(
            format=pyaudio.paInt16,
            channels=1,
            rate=44100,
            output=True,
            frames_per_buffer=1024
        )

    def connect(self):
        """连接到服务端"""
        try:
            self.sock.connect((self.host, self.port))
            self.running = True
            # 启动帧接收线程
            threading.Thread(target=self.receive_frames, daemon=True).start()
            # 启动输入处理线程
            self.input_thread = threading.Thread(target=self.send_input_events, daemon=True)
            self.input_thread.start()
            return True
        except Exception as e:
            print(f"连接错误: {e}")
            return False

    def receive_frames(self):
        """接收并处理服务端发送的音视频数据"""
        try:
            while self.running:
                # 先接收数据类型（1字节：V=视频，A=音频）
                data_type = self.sock.recv(1)
                if not data_type:
                    break

                # 接收数据长度（4字节）
                size_data = self.sock.recv(4)
                if not size_data:
                    break
                size = struct.unpack('!I', size_data)[0]

                # 接收完整数据
                frame_data = b''
                while len(frame_data) < size:
                    chunk = self.sock.recv(size - len(frame_data))
                    if not chunk:
                        break
                    frame_data += chunk

                if len(frame_data) != size:
                    continue  # 数据不完整则跳过

                # 处理视频数据
                if data_type == b'V':
                    try:
                        decompressed = zlib.decompress(frame_data)
                        with self.frame_lock:
                            self.latest_frame = decompressed
                            self.last_frame_time = time.time()
                    except Exception as e:
                        print(f"视频解压错误: {e}")

                # 处理音频数据
                elif data_type == b'A':
                    try:
                        decompressed_audio = zlib.decompress(frame_data)
                        self.audio_stream.write(decompressed_audio)
                    except Exception as e:
                        print(f"音频处理错误: {e}")

        except ConnectionResetError:
            print("服务器已断开连接")
            self.running = False
            msgb.showinfo("连接断开", "与服务器的连接已丢失")
        except Exception as e:
            print(f"接收帧错误: {e}")
        finally:
            self.running = False

    def send_input_events(self):
        """空循环维持线程，实际输入在UI事件中处理"""
        try:
            while self.running:
                time.sleep(0.01)
        except Exception as e:
            print(f"输入线程错误: {e}")

    def get_latest_frame(self):
        """获取最新视频帧"""
        with self.frame_lock:
            return self.latest_frame

    def close(self):
        """关闭连接并释放资源"""
        self.running = False
        try:
            self.sock.close()
            self.audio_stream.stop_stream()
            self.audio_stream.close()
            self.audio.terminate()
        except Exception as e:
            print(f"关闭错误: {e}")


class RemoteViewerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("远程桌面客户端")
        self.client = RemoteClient()
        self.setup_ui()
        self.setup_controls()
        self.setup_input_handling()
        self.update_ui()
        self.frame_count = 0
        self.last_fps_update = time.time()

    def setup_ui(self):
        """设置UI界面"""
        self.main_frame = tk.Frame(self.root)
        self.main_frame.pack(fill=tk.BOTH, expand=True)

        self.canvas = tk.Canvas(self.main_frame, bg='black')
        self.canvas.pack(fill=tk.BOTH, expand=True)

        self.control_frame = tk.Frame(self.root)
        self.control_frame.pack(fill=tk.X)

        self.status_label = tk.Label(
            self.control_frame,
            text="未连接",
            anchor=tk.W
        )
        self.status_label.pack(side=tk.LEFT, fill=tk.X, expand=True)

        self.fps_label = tk.Label(
            self.control_frame,
            text="FPS: 0",
            anchor=tk.E
        )
        self.fps_label.pack(side=tk.RIGHT)

    def setup_controls(self):
        """设置菜单栏控制"""
        menubar = tk.Menu(self.root)
        filemenu = tk.Menu(menubar, tearoff=0)
        filemenu.add_command(label="连接", command=self.connect_to_server)
        filemenu.add_command(label="断开", command=self.disconnect)
        filemenu.add_separator()
        filemenu.add_command(label="退出", command=self.root.quit)
        menubar.add_cascade(label="文件", menu=filemenu)
        self.root.config(menu=menubar)

    def setup_input_handling(self):
        """绑定鼠标和键盘事件"""
        self.canvas.bind("<Motion>", self.handle_mouse_move)
        self.canvas.bind("<Button-1>", lambda e: self.handle_mouse_click('L'))  # 左键
        self.canvas.bind("<Button-3>", lambda e: self.handle_mouse_click('R'))  # 右键
        self.root.bind("<Key>", self.handle_key_press)

    def handle_mouse_move(self, event):
        """处理鼠标移动事件（坐标转换为服务器屏幕比例）"""
        if self.client.running:
            try:
                # 获取客户端窗口和原始图像尺寸
                canvas_width = self.canvas.winfo_width()
                canvas_height = self.canvas.winfo_height()
                frame_data = self.client.get_latest_frame()

                if frame_data and canvas_width > 0 and canvas_height > 0:
                    # 计算缩放比例
                    img = Image.open(BytesIO(frame_data))
                    orig_width, orig_height = img.size
                    scale_x = orig_width / canvas_width
                    scale_y = orig_height / canvas_height

                    # 转换为服务器屏幕坐标
                    x = int(event.x * scale_x)
                    y = int(event.y * scale_y)
                    self.client.sock.sendall(b'M' + struct.pack('!II', x, y))
            except Exception as e:
                print(f"鼠标移动错误: {e}")

    def handle_mouse_click(self, button):
        """处理鼠标点击事件"""
        if self.client.running:
            try:
                self.client.sock.sendall(b'C' + button.encode())
            except Exception as e:
                print(f"鼠标点击错误: {e}")

    def handle_key_press(self, event):
        """处理键盘事件"""
        if self.client.running and len(event.char) == 1:
            try:
                self.client.sock.sendall(b'K' + struct.pack('!I', ord(event.char)))
            except Exception as e:
                print(f"键盘输入错误: {e}")

    def connect_to_server(self):
        """连接到服务器"""
        if not self.client.running:
            try:
                ip = strq.askstring("连接", "请输入服务器IP:", initialvalue=self.client.host)
                port = strq.askstring("连接", "请输入端口号:", initialvalue=str(self.client.port))
                if not ip or not port:
                    return
                self.client.host = ip
                self.client.port = int(port)
                if self.client.connect():
                    self.status_label.config(text=f"已连接到 {ip}:{port}")
                else:
                    msgb.showerror("错误", "连接失败")
            except Exception as e:
                msgb.showerror("错误", f"无效输入: {e}")

    def disconnect(self):
        """断开连接"""
        if self.client.running:
            self.client.close()
            self.status_label.config(text="未连接")

    def update_ui(self):
        """更新UI显示（视频帧和FPS）"""
        if self.client.running:
            frame_data = self.client.get_latest_frame()
            if frame_data:
                try:
                    # 显示视频帧
                    img = Image.open(BytesIO(frame_data))
                    self.photo = ImageTk.PhotoImage(img)
                    self.canvas.create_image(0, 0, image=self.photo, anchor=tk.NW)

                    # 计算FPS
                    self.frame_count += 1
                    current_time = time.time()
                    if current_time - self.last_fps_update >= 1.0:
                        self.fps_label.config(text=f"FPS: {self.frame_count}")
                        self.frame_count = 0
                        self.last_fps_update = current_time
                except Exception as e:
                    print(f"图像处理错误: {e}")

        # 10ms后刷新UI
        self.root.after(10, self.update_ui)


if __name__ == "__main__":
    root = tk.Tk()
    app = RemoteViewerApp(root)
    root.mainloop()